import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize
from scipy.signal import savgol_filter
import h5py
from tools.Fdecay_fhalf import Fdecay_fhalf

# 系统参数配置
class Config:
    RQ = 8893  # 电阻-品质因数乘积
    Ql = 76.7  # 外部品质因数
    beta = 1.2
    Fs = 117.402e6 / 14  # 采样率 [Hz]
    Ts = 1 / Fs  # 采样时间间隔 [s]

# 数据源配置
DATA_FILE = r'E:\MLdata\SimData\data.h5'

def read_hdf5_dataset(file_path, set_index):
    """
    从HDF5文件读取指定数据集的数据
    
    Args:
        file_path (str): HDF5文件路径
        set_index (int): 数据集索引
    
    Returns:
        dict: 包含所有电压数据的字典
    """
    with h5py.File(file_path, 'r') as f:
        set_name = f'set_{set_index}'
        if set_name not in f:
            raise ValueError(f"数据集 {set_name} 不存在")
            
        dataset = f[set_name]
        voltages = {
            "Vcav_I": dataset['Vc_real'][:],
            "Vcav_Q": dataset['Vc_imag'][:],
            "Vfor_I": dataset['Vf_star_real'][:],
            "Vfor_Q": dataset['Vf_star_imag'][:],
            "Vrefl_I": dataset['Vr_star_real'][:],
            "Vrefl_Q": dataset['Vr_star_imag'][:]
        }
        
        time = f['time'][:]
        
    return voltages, time

class VoltageProcessor:
    @staticmethod
    def clean_voltages(voltages_IQ):
        """清除电压数据中的无效值（NaN或无穷大）"""
        return {key: np.nan_to_num(value) for key, value in voltages_IQ.items()}

    @staticmethod
    def calculate_calibration_matrix(voltages_IQ):
        """计算校准矩阵和校准后的电压值"""
        voltages_IQ = VoltageProcessor.clean_voltages(voltages_IQ)
        
        # 提取IQ分量
        Vcav_I, Vcav_Q = voltages_IQ["Vcav_I"], voltages_IQ["Vcav_Q"]
        Vfor_I, Vfor_Q = voltages_IQ["Vfor_I"], voltages_IQ["Vfor_Q"]
        Vrefl_I, Vrefl_Q = voltages_IQ["Vrefl_I"], voltages_IQ["Vrefl_Q"]
        
        # 构建校准矩阵
        A_matrix = np.column_stack((
            np.append(Vfor_I, Vfor_Q),
            np.append(-Vfor_Q, Vfor_I),
            np.append(Vrefl_I, Vrefl_Q),
            np.append(-Vrefl_Q, Vrefl_I)
        ))
        y_array = np.append(Vcav_I, Vcav_Q)

        try:
            C_array = np.matmul(np.linalg.pinv(A_matrix, rcond=1e-5), y_array.transpose())
        except np.linalg.LinAlgError as e:
            print(f"SVD计算失败: {e}")
            return None

        # 计算校准后的电压值
        Vfor_cali = (Vfor_I + 1j * Vfor_Q) * (C_array[0] + 1j * C_array[1])
        Vrefl_cali = (Vrefl_I + 1j * Vrefl_Q) * (C_array[2] + 1j * C_array[3])

        return {
            'Vfor_I_cali': Vfor_cali.real,
            'Vfor_Q_cali': Vfor_cali.imag,
            'Vrefl_I_cali': Vrefl_cali.real,
            'Vrefl_Q_cali': Vrefl_cali.imag,
            'C_array': C_array
        }

    @staticmethod
    def calculate_Z(Vf, Vr):
        """计算阻抗"""
        # 基本参数设置
        Toffset = 1e-6       # 偏移时间 [us]
        R = 0.95            # 阈值比例
        Noffset = int(np.ceil(Toffset * Config.Fs))
        
        try:
            # 1. 找到RF关闭时刻
            DiffPf = np.abs(np.diff(np.abs(Vf)))
            if len(DiffPf) == 0:
                print("警告：信号太短，无法计算阻抗")
                return 0
                
            Noff = np.argmax(DiffPf) + 1
            
            # 2. 计算稳态振幅
            if Noff < 4*Noffset:
                print("警告：RF关闭时刻太早，无法计算稳态振幅")
                return 0
                
            ssPtA = np.mean(np.abs(Vr)[Noff - 4*Noffset : Noff - Noffset])
            if np.isnan(ssPtA) or ssPtA == 0:
                print("警告：稳态振幅计算失败")
                return 0
                
            # 3. 计算阈值和衰减结束点
            ssPtA_threshold = ssPtA * R
            above_threshold = np.where(np.abs(Vr) > ssPtA_threshold)[0]
            if len(above_threshold) == 0:
                print("警告：找不到衰减结束点")
                return 0
                
            Ndecay = np.max(above_threshold)
            
            # 4. 计算阻抗
            start_idx = Noff + 2
            end_idx = Noff + 10
            
            if start_idx >= len(Vf) or end_idx >= len(Vf):
                print("警告：计算区间超出信号范围")
                return 0
                
            # 避免除以零
            Vf_segment = Vf[start_idx:end_idx]
            if np.any(np.abs(Vf_segment) < 1e-10):
                print("警告：前向电压信号太弱")
                return 0
                
            Z = -Vr[start_idx:end_idx] / Vf_segment
            mean_Z = np.mean(Z)
            
            # 检查结果
            if np.isnan(mean_Z) or np.isinf(mean_Z):
                print("警告：计算得到无效阻抗值")
                return 0
                
            print(f"计算得到的阻抗Z = {mean_Z:.3f}")
            return mean_Z
            
        except Exception as e:
            print(f"计算阻抗时出错：{str(e)}")
            return 0

class OptimizationCalculator:
    @staticmethod
    def calculate_y_a(a, Z, Vc, Vf, Vr):
        """计算y_a参数"""
        if Z == 0:
            print("警告：阻抗Z为零，无法进行计算")
            return np.zeros_like(Vc)  # 返回一个与 Vc 相同形状的零数组
        
        w_half = Fdecay_fhalf(Vc, Vf, Config.Fs)
        a_complex = complex(a[0], a[1])
        Vf_cal = a_complex * Vf + a_complex / Z * Vr
        rho = 2 * np.abs((Config.beta / (1 + Config.beta)) * Vf_cal)
        return w_half * (rho * np.cos(np.angle(Vc) - np.angle(Vf_cal)) - np.abs(Vc))

    @staticmethod
    def error_function(a, Z, Vc, Vf, Vr):
        """计算误差函数"""
        # 计算导数并平滑处理
        raw_dVc_dt = np.gradient(np.abs(Vc), 1/Config.Fs)
        # 使用Savitzky-Golay滤波器进行平滑处理
        # window_length必须是奇数且大于polyorder
        dVc_dt = savgol_filter(raw_dVc_dt, window_length=51, polyorder=3)
        
        y_a = OptimizationCalculator.calculate_y_a(a, Z, Vc, Vf, Vr)
        return np.sum((dVc_dt - y_a) ** 2)

    @staticmethod
    def optimize_parameters(Vc, Vf, Vr, Z):
        """优化a参数"""
        if Z == 0:
            print("警告：阻抗Z为零，无法进行参数优化")
            return complex(0, 0)  # 返回一个默认值
        
        result = minimize(
            OptimizationCalculator.error_function,
            [0, 0.0],
            args=(Z, Vc, Vf, Vr),
            method='Nelder-Mead'
        )
        print(f"最小误差: {result.fun}")
        return complex(result.x[0], result.x[1])

class Plotter:
    @staticmethod
    def plot_results(t_us, dVc_dt, y_a_opt, a_rounded):
        """绘制导数结果图表"""
        plt.figure(figsize=(8, 6))
        
        # 绘制曲线
        plt.plot(t_us, dVc_dt, label='dVc/dt', linewidth=1.5)
        plt.plot(t_us, y_a_opt, label=f'y_a (a/X={a_rounded})', 
                linestyle='--', linewidth=1.5)

        # 设置标签和字体
        plt.xlabel(r'Time [$\mu s$]', fontname='Arial', fontsize=12)
        plt.ylabel('Differential of $ \\left | V_c \\right |  $', 
                  fontname='Arial', fontsize=12)
        
        # 设置刻度
        plt.xticks(fontname='Arial', fontsize=12)
        plt.yticks(fontname='Arial', fontsize=12)
        
        # 添加网格和图例
        plt.grid(True, linestyle='--', linewidth=0.7)
        plt.legend(loc='best', fontsize=12, frameon=True, 
                  fancybox=True, edgecolor='black', framealpha=1)
        
        # 设置边框
        ax = plt.gca()
        for spine in ax.spines.values():
            spine.set_linewidth(1.5)
            
        plt.tight_layout()
        plt.show()

    @staticmethod
    def plot_voltage_comparison(t_us, Vc_star, Vf_star, Vr_star, X, Y):
        """绘制电压比较图表
        
        Args:
            t_us: 时间数组 (微秒)
            Vc_star: 腔体电压
            Vf_star: 前向电压
            Vr_star: 反射电压
            X, Y: 校准系数
        """
        plt.figure(figsize=(10, 6))
        
        # 计算 X*Vf + Y*Vr
        V_combined = X * Vf_star + Y * Vr_star
        
        # 绘制幅值曲线
        plt.plot(t_us, np.abs(Vc_star), label='|Vc|', linewidth=1.5)
        plt.plot(t_us, np.abs(V_combined), 
                label='|X*Vf + Y*Vr|', linestyle='--', linewidth=1.5)
        
        # 设置标签和字体
        plt.xlabel(r'Time [$\mu s$]', fontname='Arial', fontsize=12)
        plt.ylabel('Voltage Magnitude', fontname='Arial', fontsize=12)
        
        # 设置刻度
        plt.xticks(fontname='Arial', fontsize=12)
        plt.yticks(fontname='Arial', fontsize=12)
        
        # 添加网格和图例
        plt.grid(True, linestyle='--', linewidth=0.7)
        plt.legend(loc='best', fontsize=12, frameon=True, 
                  fancybox=True, edgecolor='black', framealpha=1)
        
        # 设置边框
        ax = plt.gca()
        for spine in ax.spines.values():
            spine.set_linewidth(1.5)
            
        plt.tight_layout()
        plt.show()

def process_file(set_index=0):
    """处理指定索引的数据集"""
    # 读取并处理数据
    try:
        voltages_IQ, time = read_hdf5_dataset(DATA_FILE, set_index)
    except Exception as e:
        print(f"处理数据集 {set_index} 时出错: {e}")
        return
    
    # 获取校准数据
    calibration_results = VoltageProcessor.calculate_calibration_matrix(voltages_IQ)
    if calibration_results is None:
        return

    # 准备时间数组
    t_us = time * 1e6  # 转换为微秒

    # 提取复数形式的电压值
    Vf_star = voltages_IQ['Vfor_I'] + 1j * voltages_IQ['Vfor_Q']
    Vr_star = voltages_IQ['Vrefl_I'] + 1j * voltages_IQ['Vrefl_Q']
    Vc_star = voltages_IQ['Vcav_I'] + 1j * voltages_IQ['Vcav_Q']

    # 计算阻抗和优化参数
    Z = VoltageProcessor.calculate_Z(Vf_star, Vr_star)
    a_opt = OptimizationCalculator.optimize_parameters(Vc_star, Vf_star, Vr_star, Z)

    # 计算和显示结果
    C = calibration_results['C_array']  # 确保从字典中正确获取 C_array
    X = C[0] + 1j * C[1]
    Y = C[2] + 1j * C[3]
    
    a_round = a_opt / X
    a_rounded = complex(round(a_round.real, 2), round(a_round.imag, 2))
    
    print(f'a={a_opt:.3f}, b={a_opt/Z:.3f}, c={X-a_opt:.3f}, d={Y-a_opt/Z:.3f}')
    print(f'X={X:.2f}, Y={Y:.2f}, Z={Z:.2f}')

    # 计算最终结果并绘图
    raw_dVc_dt = np.gradient(np.abs(Vc_star), 1/Config.Fs)
    # 对导数进行平滑处理
    dVc_dt = savgol_filter(raw_dVc_dt, window_length=51, polyorder=3)
    
    y_a_opt = OptimizationCalculator.calculate_y_a(
        [a_opt.real, a_opt.imag], Z, Vc_star, Vf_star, Vr_star
    )
    
    # 绘制导数比较图
    Plotter.plot_results(t_us, dVc_dt, y_a_opt, a_rounded)
    
    # 绘制电压幅值比较图
    Plotter.plot_voltage_comparison(t_us, Vc_star, Vf_star, Vr_star, X, Y)

def main():
    """主函数"""
    # 处理第一个数据集
    process_file(0)

if __name__ == '__main__':
    main()
